import os
import sys
import torch
from algorithms.enhanced_lstm_crf import EnhancedLstmCRFModel

# 添加项目根目录到路径
project_root = os.path.dirname(os.path.abspath(__file__))
sys.path.append(project_root)

def analyze_model_parameters():
    """分析增强版LSTM-CRF模型的参数数量"""
    try:
        # 设置设备
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        print(f"使用设备: {device}")
        
        # 模型文件路径
        model_path = os.path.join(project_root, 'scripts', 'plw', 'enhanced_lstm_crf_model.pth')
        
        if not os.path.exists(model_path):
            print(f"模型文件不存在: {model_path}")
            # 创建一个示例模型进行分析
            print("创建示例模型进行分析...")
            model = EnhancedLstmCRFModel(
                input_dim=10,
                hidden_dim=128,
                output_dim=10,
                output_seq_length=5,
                num_layers=1,
                dropout=0.1
            )
        else:
            print(f"加载模型: {model_path}")
            
            # 加载模型
            checkpoint = torch.load(model_path, map_location=device)
            
            # 获取模型配置
            if 'model_config' in checkpoint:
                config = checkpoint['model_config']
                print(f"模型配置: {config}")
                
                # 创建模型（移除lottery_type参数）
                if 'lottery_type' in config:
                    del config['lottery_type']
                
                model = EnhancedLstmCRFModel(**config)
                
                # 加载模型权重
                try:
                    model.load_state_dict(checkpoint['model_state_dict'])
                except RuntimeError as e:
                    if "Unexpected key(s) in state_dict" in str(e):
                        print("检测到模型层数不匹配，尝试创建2层模型...")
                        config['num_layers'] = 2
                        if 'lottery_type' in config:
                            del config['lottery_type']
                        model = EnhancedLstmCRFModel(**config)
                        model.load_state_dict(checkpoint['model_state_dict'])
                    else:
                        raise e
                
                model.to(device)
                model.eval()
                print("模型加载成功")
            else:
                print("模型检查点中缺少model_config")
                return
        
        # 分析模型参数
        print("\n" + "="*50)
        print("模型参数分析")
        print("="*50)
        
        total_params = 0
        trainable_params = 0
        
        for name, parameter in model.named_parameters():
            param_count = parameter.numel()
            total_params += param_count
            if parameter.requires_grad:
                trainable_params += param_count
            
            print(f"{name}: {param_count:,} 参数")
        
        print("\n" + "-"*50)
        print(f"总参数数量: {total_params:,}")
        print(f"可训练参数数量: {trainable_params:,}")
        print(f"不可训练参数数量: {total_params - trainable_params:,}")
        
        # 分析各层参数
        print("\n各层参数详细分析:")
        print("-"*50)
        
        # LSTM层参数
        lstm_params = 0
        for name, param in model.lstm.named_parameters():
            lstm_params += param.numel()
        print(f"LSTM层参数: {lstm_params:,}")
        
        # 增强模块参数
        enhancer_params = 0
        for name, param in model.sequence_lstm_enhancer.named_parameters():
            enhancer_params += param.numel()
        print(f"增强模块参数: {enhancer_params:,}")
        
        # 全连接层参数
        fc_params = 0
        for name, param in model.fc.named_parameters():
            fc_params += param.numel()
        print(f"全连接层参数: {fc_params:,}")
        
        # CRF层参数
        crf_params = 0
        for name, param in model.crf.named_parameters():
            crf_params += param.numel()
        print(f"CRF层参数: {crf_params:,}")
        
        # Dropout层参数
        dropout_params = 0
        for name, param in model.dropout.named_parameters():
            dropout_params += param.numel()
        print(f"Dropout层参数: {dropout_params:,}")
        
        print("\n" + "-"*50)
        calculated_total = lstm_params + enhancer_params + fc_params + crf_params + dropout_params
        print(f"计算总计参数: {calculated_total:,}")
        
        # 数据集信息
        print("\n数据集信息:")
        print("-"*50)
        plw_data_file = os.path.join(project_root, 'scripts', 'plw', 'plw_history.csv')
        if os.path.exists(plw_data_file):
            import pandas as pd
            df = pd.read_csv(plw_data_file)
            print(f"排列5历史数据条数: {len(df)}")
            print(f"参数与数据比: 1:{len(df)/total_params:.6f}")
            
            if len(df) / total_params < 10:
                print("⚠️  警告: 参数数量与数据量比例较低，可能存在过拟合风险")
            elif len(df) / total_params < 20:
                print("⚠️  注意: 参数数量与数据量比例适中，需要监控过拟合情况")
            else:
                print("✅ 参数数量与数据量比例良好，过拟合风险较低")
        else:
            print("无法找到排列5历史数据文件")
            
    except Exception as e:
        print(f"分析过程中出错: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    analyze_model_parameters()